/* 参考链接:
 https://github.com/apoorvam/aws-s3-multipart-upload
*/

package main

import (
    "bytes"
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/service/s3"
    "github.com/aws/aws-sdk-go/aws/awserr"
    "os"
    "log"
    "net/http"
)

func osLoadFile(filename string) ([]byte, *os.File, string, int64) {
    file, err := os.Open(filename)
    if err != nil {
        log.Printf("err opening file: %s", err)
        return nil, nil, "", 0
    }
    defer file.Close()
    fileInfo, _ := file.Stat()
    size := fileInfo.Size()
    buffer := make([]byte, size)
    fileType := http.DetectContentType(buffer)
    file.Read(buffer)
    return buffer, file, fileType, size
}

func awsS3PreUpload(svc *s3.S3, file *os.File, filetype string, s3bucket string, path string) *s3.CreateMultipartUploadOutput {
    real_path := path + "/" + file.Name()
    input := &s3.CreateMultipartUploadInput{
        Bucket:      aws.String(s3bucket),
        Key:         aws.String(real_path),
        ContentType: aws.String(filetype),
    }

    resp, err := svc.CreateMultipartUpload(input)
    if err != nil {
        log.Println(err.Error())
        return nil
    }
    log.Println("Created multipart upload request")
    return resp
}

func awsS3MultiPartsUpload(svc *s3.S3, resp *s3.CreateMultipartUploadOutput, buffer []byte, size int64) {
    var curr, partLength int64
    var remaining = size
    var completedParts []*s3.CompletedPart
    partNumber := 1
    for curr = 0; remaining != 0; curr += partLength {
        if remaining < maxPartSize {
            partLength = remaining
        } else {
            partLength = maxPartSize
        }
        completedPart, err := uploadPart(svc, resp, buffer[curr:curr+partLength], partNumber)
        if err != nil {
            log.Println(err.Error())
            err := abortMultipartUpload(svc, resp)
            if err != nil {
                log.Println(err.Error())
            }
            return
        }
        remaining -= partLength
        partNumber++
        completedParts = append(completedParts, completedPart)
    }

    completeResponse, err := completeMultipartUpload(svc, resp, completedParts)
    if err != nil {
        log.Println(err.Error())
        return
    }

    log.Printf("Successfully uploaded file: %s\n", completeResponse.String())
}

func completeMultipartUpload(svc *s3.S3, resp *s3.CreateMultipartUploadOutput, completedParts []*s3.CompletedPart) (*s3.CompleteMultipartUploadOutput, error) {
    completeInput := &s3.CompleteMultipartUploadInput{
        Bucket:   resp.Bucket,
        Key:      resp.Key,
        UploadId: resp.UploadId,
        MultipartUpload: &s3.CompletedMultipartUpload{
            Parts: completedParts,
        },
    }
    return svc.CompleteMultipartUpload(completeInput)
}

func uploadPart(svc *s3.S3, resp *s3.CreateMultipartUploadOutput, fileBytes []byte, partNumber int) (*s3.CompletedPart, error) {
    tryNum := 1
    partInput := &s3.UploadPartInput{
        Body:          bytes.NewReader(fileBytes),
        Bucket:        resp.Bucket,
        Key:           resp.Key,
        PartNumber:    aws.Int64(int64(partNumber)),
        UploadId:      resp.UploadId,
        ContentLength: aws.Int64(int64(len(fileBytes))),
    }

    for tryNum <= maxRetries {
        uploadResult, err := svc.UploadPart(partInput)
        if err != nil {
            if tryNum == maxRetries {
                if aerr, ok := err.(awserr.Error); ok {
                    return nil, aerr
                }
                return nil, err
            }
            log.Printf("Retrying to upload part #%v\n", partNumber)
            tryNum++
        } else {
            log.Printf("Uploaded part #%v\n", partNumber)
            return &s3.CompletedPart{
                ETag:       uploadResult.ETag,
                PartNumber: aws.Int64(int64(partNumber)),
            }, nil
        }
    }
    return nil, nil
}

func abortMultipartUpload(svc *s3.S3, resp *s3.CreateMultipartUploadOutput) error {
    log.Println("Aborting multipart upload for UploadId#" + *resp.UploadId)
    abortInput := &s3.AbortMultipartUploadInput{
        Bucket:   resp.Bucket,
        Key:      resp.Key,
        UploadId: resp.UploadId,
    }
    _, err := svc.AbortMultipartUpload(abortInput)
    return err
}


func (this *S3Mgmt) MultiPartsUpload(localfile, bucket, path string) error {
    s3Svc, _ := this.NewSessionConnect()
    fileBuffer,file,fileType,size := osLoadFile(localfile)
    resp := awsS3PreUpload(s3Svc, file, fileType, bucket, path) 
    awsS3MultiPartsUpload(s3Svc, resp, fileBuffer, size)
    return nil
}